'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 9_study_cohere.py
* @Time: 2025/10/31
* @All Rights Reserve By Brtc
'''
import dotenv
import weaviate
from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank
from langchain_openai import OpenAIEmbeddings
from langchain_weaviate import WeaviateVectorStore

dotenv.load_dotenv()
#1、创建向量模型
embedding = OpenAIEmbeddings(model="text-embedding-3-small")
client = weaviate.connect_to_local("192.168.106.129", 8080)
#3、将数据存储到向量数据库中
db = WeaviateVectorStore(
    client,
    index_name="TestDemo",
    text_key="text",
    embedding=embedding
)
rerank = CohereRerank(model="rerank-multilingual-v3.0")

#3、构建重排序检索器
retriever = ContextualCompressionRetriever(
    base_retriever=db.as_retriever(),
    base_compressor=rerank
)
#4、执行搜索并排序
search_docs = retriever.invoke("关于llmops的应用配置信息有那些？")
for doc in search_docs:
    print("==================================================")
    print(doc)

client.close()
